feat(jit): forward sub-part geometry env vars to the JIT - #1
Conversation
`hybrid_dispatch_unordered.cuh` gates the sub-part geometry behind `#ifndef` (`EP_NUM_SUB_PARTS` 2, `EP_MIN_SUB_TOKENS` 1, `EP_SM100_MIN_SUB_TOKENS` 15), but nothing in the tree sets those macros, so the only way to try a different split is to edit the header and reinstall. Forward the three names as JIT `-D` flags, following the `EP_NUM_TOPK_IDX_BITS` block immediately above (and its `EP_JIT_EXTRA_FLAGS` TODO). All three are device-only -- no host translation unit reads them -- so a JIT-only define cannot desync host and device sizing. `flags` is part of `kernel_signature`, so changing the env re-JITs instead of serving a cached cubin. Unset => no behaviour change.
efe26c2 to
5118c2e
Compare
|
Rebased onto The old branch showed a 3,000-line diff over 20 files, which was an artefact, not this change. It was cut from Measurement caveat: the numbers in the description were taken on |
2 participants
Re-post of Xuan-1998/DeepEP#42 against this repo, rebased onto
main(ec623f3) and re-measured on that tree — the old PR's numbers were taken on the pre-refactor tree, so they are not quoted here.What
hybrid_dispatch_unordered.cuhgates its sub-part geometry behind#ifndef:EP_NUM_SUB_PARTSEP_MIN_SUB_TOKENSEP_SM100_MIN_SUB_TOKENSNothing in the tree ever sets those macros, so the only way to try a different split today is to edit the header and reinstall. This forwards the three names as JIT
-Dflags, following theEP_NUM_TOPK_IDX_BITSblock immediately above it (and that block'sEP_JIT_EXTRA_FLAGSTODO).Unset ⇒ byte-identical behaviour.
Why it is safe
kNumMaxTokensPerRank.)flagsis part ofkernel_signature, so changing the env re-JITs instead of silently serving a cached cubin compiled with a different geometry.Verified end-to-end
EP_JIT_PRINT_COMPILER_COMMAND=1 EP_NUM_SUB_PARTS=1puts-DEP_NUM_SUB_PARTS=1on the nvcc line, and the resulting cubin lands in a distinct JIT cache entry.Measured effect of the knob it exposes
tests/elastic/test_ep.py, 2 ×p5en.48xlarge(8×H200 + 16 EFA each), EP8×2 = 16 ranks,--hidden=7168 --num-topk=8 --num-experts=256 --num-sms=12 --allow-hybrid-mode=1 --prefer-overlap-with-compute=0 --test-first-only. 3 reps, variants interleaved within each rep, each variant on its ownEP_JIT_CACHE_DIR. Mean over all 16 ranks then over reps; ± is stdev across reps. All 48 rounds exited 0.EP_NUM_SUB_PARTS=1on its own:EP_NUM_SUB_PARTS=1So on its own the knob is roughly a wash — a small prefill win, a small decode regression. Its value is that it composes: stacked with #2 it takes decode dispatch from 367.0 µs to 166.1 ± 0.4 µs (−54.7%), where that PR alone reaches 239.8 µs (−34.7%). That is the case this PR is really enabling — being able to find such a combination without a rebuild.
This PR is pure plumbing and changes no default, so it is worth taking independently of whether #2 is accepted.